import backtrader as bt

from strategy.utils.BaoStockPandasData import BaoStockPandasData
from utils.DataSource import get_stock_data2

class SMACrossoverStrategy(bt.Strategy):
    params = (
        ('short_period', 10),  # 短期均线周期
        ('long_period', 30),  # 长期均线周期
        ('order_percentage', 0.95),  # 交易资金占总资金的百分比
    )

    def __init__(self):
        self.data_close = self.datas[0].close
        self.order = None
        self.sma_short = bt.indicators.SimpleMovingAverage(self.data_close, period=self.params.short_period)
        self.sma_long = bt.indicators.SimpleMovingAverage(self.data_close, period=self.params.long_period)

    def next(self):
        if self.order:
            return  # 如果有订单在执行，则不执行新的买卖操作

        # 当短期均线向上穿越长期均线时买入
        if self.sma_short[0] > self.sma_long[0] and not self.position:
            amount_to_invest = (self.params.order_percentage * self.broker.cash) / self.data_close[0]
            self.buy(size=amount_to_invest)
            self.log(f"BUY EXECUTED --- Price: {self.data_close[0]}")

        # 当短期均线向下穿越长期均线时卖出
        elif self.sma_short[0] < self.sma_long[0] and self.position:
            self.sell()
            self.log(f"SELL EXECUTED --- Price: {self.data_close[0]}")

    def log(self, txt, dt=None):
        ''' Logging function for this strategy'''
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return  # 订单状态未最终确定，等待

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f"BUY EXECUTED --- Price: {order.executed.price}, Cost: {order.executed.value}, Commission: {order.executed.comm}")
            else:  # Sell
                self.log(f"SELL EXECUTED --- Price: {order.executed.price}, Size: {order.executed.size}, Commission: {order.executed.comm}")

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        self.order = None

# 回测设置示例
if __name__ == '__main__':
    cerebro = bt.Cerebro()

    # 假设已经通过某种方式（如Yahoo Finance）获取了数据
    df = get_stock_data2("sh.600000", "2020-01-01", "2023-12-31")
    data = BaoStockPandasData(dataname=df)
    cerebro.adddata(data)

    cerebro.addstrategy(SMACrossoverStrategy, short_period=10, long_period=30)

    cerebro.broker.setcash(100000.0)
    cerebro.broker.set_coc(True)  # 现金交易模式

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    cerebro.run()
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    
    cerebro.plot()

    
#在Backtrader中实现均线收敛发散（Moving Average Convergence Divergence, MACD）策略，实际上是指利用MACD指标来进行交易决策，而非直接基于简单移动平均线（SMA）或指数移动平均线（EMA）的直接收敛与发散现象。然而，基于您的要求，我们可以构建一个使用SMA或EMA收敛与发散概念的策略，例如利用短期和长期均线的交叉来模拟类似MACD的交易逻辑。以下是一个基于SMA收敛发散思想的简单策略示例：
#这个策略中，我们定义了短期（short_period）和长期（long_period）均线周期，当短期均线向上穿越长期均线时，视为买入信号；当短期均线向下穿越长期均线时，视为卖出信号。这实质上是利用均线收敛（交叉）作为市场趋势变化的信号，类似于MACD策略中利用快线和慢线的交叉来判断买卖时机。